package com.sheepone.util.weka;

import com.sheepone.util.algorithm.analysis.DoubleAnalysisUtil;
import weka.clusterers.SimpleKMeans;
import weka.core.DistanceFunction;
import weka.core.EuclideanDistance;
import weka.core.Instance;
import weka.core.Instances;

import java.util.Arrays;

/**
 * @author Miss.杨
 * @description
 * @since 2024/3/31 - 11:12
 */
public class KMeansUtil {

    /**
     * 计算样本i到其所在簇内其他样本的平均距离
     *
     * @param kmeansResult
     * @param clusterId
     * @param instanceIndex
     * @return
     */
    private static double calculateIntraClusterDistance(KmeansResult kmeansResult, int clusterId, int instanceIndex) {
        Instances data = kmeansResult.getInstances();
        int[] clusterAssignments = kmeansResult.getClusterAssignments();
        double sumDist = 0.0;
        int count = 0;
        Instance instance = data.instance(instanceIndex);
        for (int i = 0; i < data.numInstances(); i++) {
            if (clusterAssignments[i] == clusterId && instanceIndex != i) {
                sumDist += WekaUtil.calculateEuclideanDistance(instance, data.instance(i));
                count++;
            }
        }
        return count > 0 ? sumDist / count : Double.POSITIVE_INFINITY;
    }


    /**
     * 计算样本i到最近邻簇的平均距离（这里假设所有簇都有至少一个样本）
     *
     * @param kmeansResult
     * @param clusterId
     * @param instanceIndex
     * @return
     */
    private static double calculateNearestClusterDistance(KmeansResult kmeansResult, int clusterId, int instanceIndex) {
        Instances data = kmeansResult.getInstances();
        int[] clusterAssignments = kmeansResult.getClusterAssignments();
        double minDist = Double.MAX_VALUE;
        Instance instance = data.instance(instanceIndex);
        for (int c = 0; c < kmeansResult.getClusterCentroids().size(); c++) {
            if (c != clusterId) {
                double clusterDist = 0.0;
                int count = 0;
                for (int i = 0; i < data.numInstances(); i++) {
                    if (clusterAssignments[i] == c) {
                        clusterDist += WekaUtil.calculateEuclideanDistance(instance, data.instance(i));
                        count++;
                    }
                }
                clusterDist /= count;
                minDist = Math.min(minDist, clusterDist);
            }
        }
        return minDist;
    }


    /**
     * 计算轮廓系数
     *
     * @param kmeansResult
     * @param instanceIndex
     * @return
     */
    private static double calculateSilhouette(KmeansResult kmeansResult, int instanceIndex) {
        int[] clusterAssignments = kmeansResult.getClusterAssignments();
        int clusterId = clusterAssignments[instanceIndex];
        double a = calculateIntraClusterDistance(kmeansResult, clusterId, instanceIndex);
        double b = calculateNearestClusterDistance(kmeansResult, clusterId, instanceIndex);

        // 轮廓系数为(b - a) / max(a, b)
        return (b - a) / Math.max(a, b);
    }


    /**
     * 计算平均轮廓系数
     *
     * @param kmeansResult
     * @return
     */
    private static double calculateMeanSilhouette(KmeansResult kmeansResult) {
        double sum = 0.0;
        Instances data = kmeansResult.getInstances();
        int dataSize = data.numInstances();
        for (int i = 0; i < dataSize; i++) {
            sum += calculateSilhouette(kmeansResult, i);
        }
        return sum / dataSize;
    }


    /**
     * 计算最佳聚类数
     *
     * @param data
     * @return
     * @throws Exception
     */
    public static KmeansResult bestKmeans(double[][] data) {
        Instances instances = WekaUtil.convertToInstances(data);
        return bestKmeans(instances);
    }

    /**
     * 计算最佳聚类数
     *
     * @param data
     * @return
     * @throws Exception
     */
    public static KmeansResult bestKmeans(Instances data) {
        KmeansResult[] results = new KmeansResult[9];
        // 在区间[2,10]选择出最佳聚类数，因为weka的Kmeans算法允许聚类数小于预期k的情况所以这里多往后尝试三个数
        double[] silhouette = new double[9];
        Arrays.fill(silhouette, DoubleAnalysisUtil.DOUBLE_MIN_VALUE);

        for (int k = 2; k <= 13; k++) {
            KmeansResult kmeansResult = kmeans(k, data);
            int index = kmeansResult.getClusterCentroids().size() - 2;
            if (index >=0  && index <= 8 && results[index] == null) {
                // 计算平均轮廓系数
                silhouette[index] = calculateMeanSilhouette(kmeansResult);
                results[index] = kmeansResult;
            }
        }
        int index = DoubleAnalysisUtil.maxValueIndex(silhouette);
        return results[index];
    }

    /**
     * KMeans算法，默认距离函数为 EuclideanDistance
     * 与matlab算法有所差异
     * 聚类数小于k的原因是聚类过程中出现空簇如果没有特殊的处理，这个问题是无法在后续的循环中自行解决的。
     * weka的实现是允许输出聚类数小于k的，出现空簇就直接将空簇删除;
     * 而matlab的km有一个参数用于解决聚类数小于k的问题，例如直接选择最远离聚类中心的点作为一个新的簇。
     *
     * @param clusterNum 目标聚类数
     * @param instances  参与聚类的原始数据
     * @return
     * @throws Exception
     */
    public static KmeansResult kmeans(int clusterNum, Instances instances) {
        return kmeans(clusterNum, instances, new EuclideanDistance());
    }

    /**
     * KMeans算法
     * 与matlab算法有所差异
     * 聚类数小于k的原因是聚类过程中出现空簇如果没有特殊的处理，这个问题是无法在后续的循环中自行解决的。
     * weka的实现是允许输出聚类数小于k的，出现空簇就直接将空簇删除;
     * 而matlab的km有一个参数用于解决聚类数小于k的问题，例如直接选择最远离聚类中心的点作为一个新的簇。
     *
     * @param clusterNum       目标聚类数
     * @param data             参与聚类的原始数据
     * @param distanceFunction 距离函数
     * @return
     */
    public static KmeansResult kmeans(int clusterNum, Instances data, DistanceFunction distanceFunction) {
        SimpleKMeans kmeans = new SimpleKMeans();
        try {
            // 期望的簇点数量
            kmeans.setNumClusters(clusterNum);
            // 距离算法
            kmeans.setDistanceFunction(distanceFunction);
            // 构建数据集模型
            kmeans.buildClusterer(data);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        // 解析结果
        return new KmeansResult(kmeans, data);
    }
}
